{
  "nbformat": 4,
  "nbformat_minor": 0,
  "metadata": {
    "colab": {
      "name": "autoestimator_pytorch_lenet_mnist.ipynb",
      "provenance": [],
      "toc_visible": true,
      "include_colab_link": true
    },
    "kernelspec": {
      "display_name": "Python 3",
      "language": "python",
      "name": "python3"
    },
    "language_info": {
      "codemirror_mode": {
        "name": "ipython",
        "version": 3
      },
      "file_extension": ".py",
      "mimetype": "text/x-python",
      "name": "python",
      "nbconvert_exporter": "python",
      "pygments_lexer": "ipython3",
      "version": "3.7.7"
    }
  },
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "view-in-github",
        "colab_type": "text"
      },
      "source": [
        "[![Open In Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/intel-analytics/BigDL/blob/main/python/orca/colab-notebook/quickstart/autoestimator_pytorch_lenet_mnist.ipynb)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "I1fvwI1q13mC"
      },
      "source": [
        "![image.png]()\n",
        "---"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "pO3ksbTR18JZ"
      },
      "source": [
        "##### Copyright 2016 The BigDL Authors."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "_HZJ3OR71u23"
      },
      "source": [
        "#@title Licensed under the Apache License, Version 2.0 (the \"License\");\n",
        "# you may not use this file except in compliance with the License.\n",
        "# You may obtain a copy of the License at\n",
        "#\n",
        "#     http://www.apache.org/licenses/LICENSE-2.0\n",
        "#\n",
        "# Unless required by applicable law or agreed to in writing, software\n",
        "# distributed under the License is distributed on an \"AS IS\" BASIS,\n",
        "# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n",
        "# See the License for the specific language governing permissions and\n",
        "# limitations under the License.\n",
        "#"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "tNTssjdw2Bpi"
      },
      "source": [
        "## **Environment Preparation**"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "fOosIv3t2Fhp"
      },
      "source": [
        "**Install Java 8**\n",
        "\n",
        "Run the cell on the **Google Colab** to install jdk 1.8.\n",
        "\n",
        "**Note:** if you run this notebook on your computer, root permission is required when running the cell to install Java 8. (You may ignore this cell if Java 8 has already been set up in your computer)."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "IBs-RL5p2Ia2"
      },
      "source": [
        "# Install jdk8\n",
        "!apt-get install openjdk-8-jdk-headless -qq > /dev/null\n",
        "import os\n",
        "# Set environment variable JAVA_HOME.\n",
        "os.environ[\"JAVA_HOME\"] = \"/usr/lib/jvm/java-8-openjdk-amd64\"\n",
        "!update-alternatives --set java /usr/lib/jvm/java-8-openjdk-amd64/jre/bin/java\n",
        "!java -version"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "wS20JSbY2LCJ"
      },
      "source": [
        "**Install BigDL Orca**\n",
        "\n",
        " "
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "5C1GYEhGIlWu"
      },
      "source": [
        "[Conda](https://docs.conda.io/projects/conda/en/latest/user-guide/install/) is needed to prepare the Python environment for running this example. \n",
        "\n",
        "**Note**: The following code cell is specific for setting up conda environment on Colab; for general conda installation, please refer to the [install guide](https://docs.conda.io/projects/conda/en/latest/user-guide/install/) for more details."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "S3UweucigGaR"
      },
      "source": [
        "import sys\n",
        "\n",
        "# Set current python version\n",
        "python_version = \"3.7.10\""
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "wouustbSJS2r"
      },
      "source": [
        "# Install Miniconda\n",
        "!wget https://repo.continuum.io/miniconda/Miniconda3-4.5.4-Linux-x86_64.sh\n",
        "!chmod +x Miniconda3-4.5.4-Linux-x86_64.sh\n",
        "!./Miniconda3-4.5.4-Linux-x86_64.sh -b -f -p /usr/local\n",
        "\n",
        "# Update Conda\n",
        "!conda install --channel defaults conda python=$python_version --yes\n",
        "!conda update --channel defaults --all --yes\n",
        "\n",
        "# Append to the sys.path\n",
        "_ = (sys.path\n",
        "        .append(f\"/usr/local/lib/python3.7/site-packages\"))\n",
        "\n",
        "os.environ['PYTHONHOME']=\"/usr/local\""
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "1LOrK0lQHhrh"
      },
      "source": [
        "You can install the latest pre-release version using `pip install --pre --upgrade bigdl-orca[automl]`."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "A8lgWGhG2Oij"
      },
      "source": [
        "# Install latest pre-release version of BigDL Orca \n",
        "# Installing BigDL Orca from pip will automatically install pyspark, bigdl, and their dependencies.\n",
        "!pip install --pre --upgrade bigdl-orca[automl]"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "3EFkPl4I2RJG"
      },
      "source": [
        "# Install python dependencies\n",
        "!pip install torch==1.8.1 torchvision==0.9.1"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "ATJ4YPAS2TQm"
      },
      "source": [
        "## **Automated hyper-parameter search for PyTorch using Orca APIs**\n",
        "\n",
        "In this guide we will describe how to enable automated hyper-parameter search for PyTorch using Orca `AutoEstimator` in 5 simple steps."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "LXeY_v2S24bN"
      },
      "source": [
        "# import necesary libraries and modules\n",
        "from __future__ import print_function\n",
        "import os\n",
        "import argparse\n",
        "\n",
        "from bigdl.orca import init_orca_context, stop_orca_context\n",
        "from bigdl.orca import OrcaContext"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "dn6xcsq43FRQ"
      },
      "source": [
        "### **Step 1: Init Orca Context**"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "mXMGhKPk3GYN"
      },
      "source": [
        "# recommended to set it to True when running BigDL in Jupyter notebook. \n",
        "OrcaContext.log_output = True # (this will display terminal's stdout and stderr in the Jupyter notebook).\n",
        "\n",
        "cluster_mode = \"local\"\n",
        "\n",
        "if cluster_mode == \"local\":\n",
        "    init_orca_context(cores=4, memory=\"2g\", init_ray_on_spark=True) # run in local mode\n",
        "elif cluster_mode == \"k8s\":\n",
        "    init_orca_context(cluster_mode=\"k8s\", num_nodes=2, cores=4, init_ray_on_spark=True) # run on K8s cluster\n",
        "elif cluster_mode == \"yarn\":\n",
        "    init_orca_context(\n",
        "      cluster_mode=\"yarn-client\", cores=4, num_nodes=2, memory=\"2g\", init_ray_on_spark=True, \n",
        "      driver_memory=\"10g\", driver_cores=1) # run on Hadoop YARN cluster"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "hbu_llz48oNL"
      },
      "source": [
        "This is the only place where you need to specify local or distributed mode. View [Orca Context](https://bigdl.readthedocs.io/en/latest/doc/Orca/Overview/orca-context.html) for more details.\n",
        "\n",
        "**Note**: You should export HADOOP_CONF_DIR=/path/to/hadoop/conf/dir when you run on Hadoop YARN cluster."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "RBzXvR3LNxBP"
      },
      "source": [
        "### **Step 2: Define the Model**\n",
        "You may define your model, loss and optimizer in the same way as in any standard PyTorch program."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "bZASG0afNWxd"
      },
      "source": [
        "import torch\n",
        "import torch.nn as nn\n",
        "import torch.nn.functional as F\n",
        "\n",
        "class LeNet(nn.Module):\n",
        "    def __init__(self, fc1_hidden_size=500):\n",
        "        super(LeNet, self).__init__()\n",
        "        self.conv1 = nn.Conv2d(1, 20, 5, 1)\n",
        "        self.conv2 = nn.Conv2d(20, 50, 5, 1)\n",
        "        self.fc1 = nn.Linear(4*4*50, fc1_hidden_size)\n",
        "        self.fc2 = nn.Linear(fc1_hidden_size, 10)\n",
        "\n",
        "    def forward(self, x):\n",
        "        x = F.relu(self.conv1(x))\n",
        "        x = F.max_pool2d(x, 2, 2)\n",
        "        x = F.relu(self.conv2(x))\n",
        "        x = F.max_pool2d(x, 2, 2)\n",
        "        x = x.view(-1, 4*4*50)\n",
        "        x = F.relu(self.fc1(x))\n",
        "        x = self.fc2(x)\n",
        "        return F.log_softmax(x, dim=1)\n",
        "\n",
        "criterion = nn.NLLLoss()"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "nQeViMnCgGaV"
      },
      "source": [
        "After defining your model, you need to define a *Model Creator Function* that returns an instance of your model, and a *Optimizer Creator Function* that returns a PyTorch optimizer. Note that both the *Model Creator Function* and the *Optimizer Creator Function* should take `config` as input and get the hyper-parameter values from `config`."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "hGo38BIvgGaV"
      },
      "source": [
        "def model_creator(config):\n",
        "    model = LeNet(fc1_hidden_size=config[\"fc1_hidden_size\"])\n",
        "    return model\n",
        "\n",
        "def optim_creator(model, config):\n",
        "    return torch.optim.Adam(model.parameters(), lr=config[\"lr\"])"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "o-SeParPNtK8"
      },
      "source": [
        "### **Step 3: Define Dataset**\n",
        "\n",
        "You can define the train and validation datasets using *Data Creator Function* that has one parameter of `config` and returns a PyTorch `DataLoader`."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "tQMxn902OBoR"
      },
      "source": [
        "import torch\n",
        "from torchvision import datasets, transforms\n",
        "\n",
        "torch.manual_seed(0)\n",
        "dir = '/tmp/dataset'\n",
        "test_batch_size = 640\n",
        "\n",
        "def train_loader_creator(config):\n",
        "    train_loader = torch.utils.data.DataLoader(\n",
        "        datasets.MNIST(dir, train=True, download=True,\n",
        "                       transform=transforms.Compose([\n",
        "                           transforms.ToTensor(),\n",
        "                           transforms.Normalize((0.1307,), (0.3081,))\n",
        "                       ])),\n",
        "        batch_size=config[\"batch_size\"], shuffle=True)\n",
        "    return train_loader\n",
        "\n",
        "def test_loader_creator(config):\n",
        "    test_loader = torch.utils.data.DataLoader(\n",
        "        datasets.MNIST(dir, train=False, download=True,\n",
        "                       transform=transforms.Compose([\n",
        "                           transforms.ToTensor(),\n",
        "                           transforms.Normalize((0.1307,), (0.3081,))\n",
        "                       ])),\n",
        "        batch_size=test_batch_size, shuffle=False)\n",
        "    return test_loader"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "Kktbp2PHg-h7"
      },
      "source": [
        "### **Step 4: Define search space**\n",
        "\n",
        "You should define a dictionary as your hyper-parameter search space. The keys are hyper-parameter names which should be the same with those in your creators, and you can specify how you want to sample each hyper-parameter in the values of the search space. See [automl.hp](https://bigdl.readthedocs.io/en/latest/doc/PythonAPI/AutoML/automl.html#orca-automl-hp) for more details.\n"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "0pXzKUfhkEea"
      },
      "source": [
        "from bigdl.orca.automl import hp\n",
        "\n",
        "search_space = {\n",
        "    \"fc1_hidden_size\": hp.choice([500, 600]),\n",
        "    \"lr\": hp.choice([0.001, 0.003]),\n",
        "    \"batch_size\": hp.choice([160, 320, 640]),\n",
        "}"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "vLSjmlkROzPr"
      },
      "source": [
        "### **Step 5: Automatically fit and search *with* Orca AutoEstimator**"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "49Rq_fo2O34O"
      },
      "source": [
        "First, create an AutoEstimator. You can refer to [AutoEstimator API doc](https://bigdl.readthedocs.io/en/latest/doc/PythonAPI/AutoML/automl.html#orca-automl-auto-estimator) for more details."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "xskLU95rmH14"
      },
      "source": [
        "from bigdl.orca.automl.auto_estimator import AutoEstimator\n",
        "\n",
        "auto_est = AutoEstimator.from_torch(model_creator=model_creator,\n",
        "                                    optimizer=optim_creator,\n",
        "                                    loss=criterion,\n",
        "                                    logs_dir=\"/tmp/orca_automl_logs\",\n",
        "                                    resources_per_trial={\"cpu\": 2},\n",
        "                                    name=\"lenet_mnist\")"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "6PEmVuGnPFBk"
      },
      "source": [
        "Next, use the auto estimator to fit and search for the best hyper-parameter set."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "WeCoGJrGPH5M"
      },
      "source": [
        "auto_est.fit(data=train_loader_creator,\n",
        "             validation_data=test_loader_creator,\n",
        "             search_space=search_space,\n",
        "             n_sampling=2,\n",
        "             epochs=1,\n",
        "             metric=\"accuracy\")"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "tFQCQwH0UBF1"
      },
      "source": [
        "Finally, you can get the best learned model and the best hyper-parameters."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "_D1lGpM4Npi-"
      },
      "source": [
        "best_model = auto_est.get_best_model()"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "vwdULWwcOpsx"
      },
      "source": [
        "best_config = auto_est.get_best_config()\n",
        "print(best_config)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "5IAl7L57b396"
      },
      "source": [
        "You can use the best learned model and the best hyper-parameters as you want. Here, we demonstrate how to evaluate on the test dataset."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "xONG4fRNRQZJ"
      },
      "source": [
        "test_loader = test_loader_creator(best_config)\n",
        "best_model.eval()\n",
        "correct = 0\n",
        "with torch.no_grad():\n",
        "  for data, target in test_loader:\n",
        "    output = best_model(data)\n",
        "    pred = output.data.max(1, keepdim=True)[1]\n",
        "    correct += pred.eq(target.data.view_as(pred)).sum().numpy()\n",
        "accuracy = 100. * correct / len(test_loader.dataset)\n",
        "print(f\"accuracy is {accuracy}%\")"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "YASrq-VzXdpJ"
      },
      "source": [
        "You can find the accuracy of the best model has reached 98%."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "uSSReIykRPIB"
      },
      "source": [
        "# stop orca context when program finishes\n",
        "stop_orca_context()"
      ],
      "execution_count": null,
      "outputs": []
    }
  ]
}
